import numpy as np
import datetime

from matplotlib import pyplot as plt
from matplotlib.dates import DateFormatter
from matplotlib import rcParams

config = {
    "font.family": 'serif',
    "font.size": 15,
    "mathtext.fontset": 'stix',
    "font.serif": ['SimSun'],
}
rcParams.update(config)

# 导入金价数据
import gold

prices = gold.prices

# 将金价数据转换为字典，便于后续使用
price_dict = {datetime.datetime.strptime(date, "%Y-%m-%d"): price for date, price in prices}


# 定义计算投资总额和累计收益的函数
def calculate_investment(start_date, end_date):
    start_date = datetime.datetime.strptime(start_date, "%Y-%m-%d")
    end_date = datetime.datetime.strptime(end_date, "%Y-%m-%d")

    # 生成日期范围数组，包含从开始买入日到卖出日的所有日期 np.array(): 创建一个NumPy数组。
    date_range = np.array([start_date + datetime.timedelta(days=i) for i in range((end_date - start_date).days + 1)])

    # 获取有效交易日期，即在金价数据中存在的日期
    valid_dates = np.array([date for date in date_range if date in price_dict])

    # 根据有效交易日期获取相应的金价数组
    prices_array = np.array([price_dict[date] for date in valid_dates])

    # 计算每天购买的黄金数量
    amounts = 1000 / prices_array

    # 计算投资总额
    total_investment = len(valid_dates) * 1000

    # 计算持有的黄金总量 np.sum(): 计算数组元素的总和。
    gold_amount = np.sum(amounts)

    # 计算累计收益
    total_money = gold_amount * price_dict[end_date] - total_investment

    return total_investment, total_money, valid_dates, prices_array


# 设置起始日期和结束日期
start_date = '2016-01-04'
end_date = '2016-01-29'

# 调用函数计算投资总额和累计收益
total_investment, total_money, valid_dates, prices_array = calculate_investment(start_date, end_date)
print('投资总额:', total_investment)
print('累计收益:', total_money)

# 计算每天的累计收益 np.cumsum(): 计算数组元素的累积和。
day_money = np.cumsum(1000 / prices_array * prices_array[-1] - 1000)

# plt.figure(): 创建一个新的图形。
# plt.plot(): 绘制折线图。
# plt.xlabel(): 设置x轴标签。
# plt.ylabel(): 设置y轴标签。
# plt.title(): 设置图表标题。
# plt.xticks(): 设置x轴刻度标签的旋转角度。
# plt.gca(): 获取当前坐标轴。
# plt.gca().xaxis.set_major_formatter(): 设置x轴主要刻度标签的格式。
# plt.grid(): 添加网格线到图表中。
# plt.text(): 在图表中的指定位置添加文本标签。
# plt.legend(): 添加图例到图表中，用于标识不同曲线的含义。
# plt.show(): 显示图形。

# 绘制可视化结果
plt.figure(figsize=(24, 10))
plt.plot(valid_dates, day_money, marker='o', linestyle='-', label='每日累计收益')
plt.xlabel('日期')
plt.ylabel('累计收益')
plt.title('黄金每日累计收益')
plt.xticks(rotation=30)
plt.gca().xaxis.set_major_formatter(DateFormatter('%Y-%m-%d'))  # 设置日期格式为年-月-日
plt.grid(True)

# 在每个点处显示累计收益的数值
for i in range(len(valid_dates)):
    plt.text(valid_dates[i], day_money[i], '{:.2f}'.format(day_money[i]), ha='center', va='bottom')

plt.legend()
plt.show()
